
from torch.nn import Module
import torch


class new_GNN_Layer(Module):


    def __init__(self, inF, outF, is_use_cuda=True):

        super(new_GNN_Layer, self).__init__()
        self.is_user_cuda = is_use_cuda
        self.inF = inF
        self.outF = outF

        self.linear = torch.nn.Linear(in_features=inF, out_features=outF)
        self.linear1 = torch.nn.Linear(in_features=inF, out_features=outF)

        self.interActTransform = torch.nn.Linear(in_features=inF, out_features=outF)
        self.interActTransform1 = torch.nn.Linear(in_features=inF, out_features=outF)

    def forward(self, Laplacian_mat, selfLoop, UI_Laplacian_mat, features):

        L1 = Laplacian_mat + selfLoop
        if self.is_user_cuda:
            L2 = Laplacian_mat.cuda()
            L1 = L1.cuda()
            L3 = UI_Laplacian_mat.cuda()
        else:
            L2 = Laplacian_mat
            L3 = UI_Laplacian_mat
        inter_feature = torch.mul(features, features)
        inter_part1 = self.linear(torch.sparse.mm(L1, features))
        #  inter_partself=self.linear(torch.sparse.mm(selfLoop.cuda(),features))
        inter_part2 = self.interActTransform(torch.sparse.mm(L2, inter_feature))

        inter_part_w1 = self.linear1(torch.sparse.mm(L3, features))
        inter_part_w2 = self.interActTransform1(torch.sparse.mm(L3, inter_feature))

        return inter_part1 + inter_part_w1 + inter_part2 + inter_part_w2
